package com.aizuda.easyManagerTool.util;

import cn.hutool.core.net.NetUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.ssh.JschUtil;
import com.aizuda.easyManagerTool.domain.dto.dbc.db.ForwardInfo;
import com.aizuda.easyManagerTool.domain.dto.terminal.SSHInfoDTO;
import com.aizuda.easyManagerTool.domain.dto.terminal.SSHWindowsDTO;
import com.aizuda.easyManagerTool.service.terminal.impl.TerminalSessionManager;
import com.jcraft.jsch.*;
import org.apache.commons.lang3.StringUtils;

import java.util.function.Consumer;

public class SSHManagerUtil {

    public static Session setSession(TerminalSessionManager.Manager manager, SSHInfoDTO info) throws Exception {
        Session session = manager.getSession();
        if(ObjectUtil.isEmpty(session)){
            if(ObjectUtil.isNotEmpty(info.getProxyId())){
                if(info.getProxyMode().equals("http")){
                    ProxyHTTP proxyHTTP = new ProxyHTTP(info.getProxyHost(),info.getProxyPort());
                    proxyHTTP.setUserPasswd(info.getProxyAccount(),info.getProxyPassword());
                    session.setProxy(proxyHTTP);
                }
                if(info.getProxyMode().equals("socket5")){
                    ProxySOCKS5 proxySOCKS5 = new ProxySOCKS5(info.getProxyHost(),info.getProxyPort());
                    proxySOCKS5.setUserPasswd(info.getProxyAccount(),info.getProxyPassword());
                    session.setProxy(proxySOCKS5);
                }
            }
            if(ObjectUtil.isNotEmpty(info.getPkId())){
                byte[] pkPassword = null;
                if(StrUtil.isNotEmpty(info.getPkPassword())){
                    pkPassword = info.getPkPassword().getBytes();
                }
                if(StrUtil.isNotBlank(info.getPkPath())) {
                    session = JschUtil.createSession(info.getServerIp(), info.getServerPort(), info.getServerAccount(), info.getPkPath(), pkPassword);
                }else{
                    session = JschUtil.createSession(info.getServerIp(), info.getServerPort(), info.getServerAccount(), info.getPkContent().getBytes(), pkPassword);
                }
            }else {
                session = JschUtil.openSession(info.getServerIp(), info.getServerPort(), info.getServerAccount(), info.getServerPassword(), info.getServerTimeout());
            }
            // 严格的主机密钥检查
            session.setConfig("StrictHostKeyChecking", "no");
        }
        if(!session.isConnected()) {
            session.connect(info.getServerTimeout());
        }
        // 持续发送心跳
        session.sendKeepAliveMsg();
        manager.setSession(session);
        return session;
    }

    /**
     * 目前仅用于数据库隧道连接
     * @param manager
     * @param info
     * @return
     * @throws Exception
     */
    public static Session forwardSession(TerminalSessionManager.Manager manager, SSHInfoDTO info) throws Exception {
        Session session = setSession(manager, info);
        ForwardInfo forwardInfo = info.getForwardInfo();
        if(ObjectUtil.isEmpty(session) || ObjectUtil.isEmpty(forwardInfo.getRPort()) || ObjectUtil.isEmpty(forwardInfo.getRHost())){
            return session;
        }
        try {
            int localPort = StringUtils.isNotBlank(forwardInfo.getLocalPort()) ? Integer.parseInt(forwardInfo.getLocalPort()) : NetUtil.getUsableLocalPort();
            forwardInfo.setLocalPort(String.valueOf(localPort));
            session.setPortForwardingL(localPort, forwardInfo.getRHost(), Integer.parseInt(forwardInfo.getRPort()));
        } catch (Exception e) {
            if (session != null && session.isConnected()) {
                session.disconnect();
            }
            throw e;
        }
        return session;
    }

    /**
     * 用于跳板机
     * @param manager
     * @param info
     * @return
     * @throws Exception
     */
    public static Session springboardSession(TerminalSessionManager.Manager manager, SSHInfoDTO info) throws Exception {
        Session session = setSession(manager, info);
        ForwardInfo forwardInfo = info.getForwardInfo();
        if(ObjectUtil.isEmpty(session) || ObjectUtil.isEmpty(forwardInfo.getRPort()) || ObjectUtil.isEmpty(forwardInfo.getRHost())){
            return session;
        }
        try {
            int localPort = !StringUtils.isBlank(forwardInfo.getLocalPort()) ? Integer.parseInt(forwardInfo.getLocalPort()) : NetUtil.getUsableLocalPort();
            forwardInfo.setLocalPort(String.valueOf(localPort));
            session.setPortForwardingL(localPort, forwardInfo.getRHost(),
                    Integer.parseInt(forwardInfo.getRPort()));
        } catch (Exception e) {
            if (session != null && session.isConnected()) {
                session.disconnect();
            }
            throw e;
        }
        return session;
    }


    public static ChannelShell setChannelShell(Session session, TerminalSessionManager.Manager manager, SSHWindowsDTO window) throws JSchException {
        //开启 shell 通道
        ChannelShell channelShell = manager.getChannelShell();
        //通道连接 超时时间3s
        if(ObjectUtil.isEmpty(channelShell)){
            channelShell = (ChannelShell) session.openChannel("shell");
            // 这里设置 vim 配色
            channelShell.setPty(true);
            if(ObjectUtil.isNotEmpty(window)) {
                channelShell.setPtyType("xterm", window.getClos(), window.getRows(), window.getWidth(), window.getHeight());
            }
        }
        if(!channelShell.isConnected()) {
            channelShell.connect(manager.getSshMessageDTO().getInfo().getServerTimeout());
        }
        manager.setChannelShell(channelShell);
        return channelShell;
    }

    /**
     * 特殊性，必须在 connect 之前设置命令
     * @param session
     * @param manager
     * @return
     * @throws JSchException
     */
    public static ChannelExec setChannelExec(Session session, TerminalSessionManager.Manager manager,String command) throws JSchException {
        //开启 shell 通道
        ChannelExec channelExec = manager.getChannelExec();
        //通道连接 超时时间3s
        channelExec = (ChannelExec) session.openChannel("exec");
        channelExec.setPty(true);
        channelExec.setCommand(command);
        channelExec.connect(session.getTimeout());
        manager.setChannelExec(channelExec);
        return channelExec;
    }


    public static ChannelSftp setChannelSFTP(Session session, TerminalSessionManager.Manager manager) throws JSchException {
        //开启 shell 通道
        ChannelSftp channelSftp = manager.getChannelSftp();
        //通道连接 超时时间3s
        if(ObjectUtil.isEmpty(channelSftp)){
            channelSftp = (ChannelSftp) session.openChannel("sftp");
        }
        if(!channelSftp.isConnected()) {
            channelSftp.connect(manager.getSshMessageDTO().getInfo().getServerTimeout());
        }
        manager.setChannelSftp(channelSftp);
        return channelSftp;
    }

    public static void templateChannelShell(TerminalSessionManager.Manager manager, Consumer<ChannelShell> consumer) throws Exception {
        Session session;
        ChannelShell channelShell;
        session = setSession(manager,manager.getSshMessageDTO().getInfo());
        channelShell = setChannelShell(session, manager,manager.getSshMessageDTO().getWindow());
        consumer.accept(channelShell);
    }

    public static void templateChannelSFTP(TerminalSessionManager.Manager manager, Consumer<ChannelSftp> consumer) throws Exception {
        Session session;
        ChannelSftp channelSftp;
        session = setSession(manager,manager.getSshMessageDTO().getInfo());
        channelSftp = setChannelSFTP(session, manager);
        consumer.accept(channelSftp);

    }


}
